#include "precomp.h"
#include <rtk/luartk_variant.h>

// 辅助模板类：将成员函数转换为静态 C 函数
#define LuaMethodWrapper(T, M) [](lua_State* L) { \
	T* me = T::self(L, 1);                        \
	return me->M(L);                              \
}

#define INST_FUNC(T, F)            \
	{                              \
		#F, LuaMethodWrapper(T, F) \
	}
#define STATIC_FUNC(T, F) \
	{                     \
		#F, T::F          \
	}

template <typename T>
struct LuaObject {
	static int Register(lua_State* L)
	{
		// 创建元表
		luaL_newmetatable(L, T::meta_name);
		luaL_setfuncs(L, T::meta_funcs, 0);
		lua_pop(L, 1);
		lua_pushstring(L, T::class_name);
		lua_newtable(L);
		luaL_setfuncs(L, T::static_funcs, 0);
		// register to gamex namespace
		lua_rawset(L, -3);
		return 0;
	}
	static int create(lua_State* L)
	{
		T* obj = (T*)lua_newuserdata(L, sizeof(T));
		if (!obj) {
			luaL_error(L, "new obj accessor failed");
		}
		luaL_getmetatable(L, T::meta_name);
		lua_setmetatable(L, -2);
		// stack contains creation arguments plus the new object itself(on top)
		new (obj) T(L);
		return 1;
	}
	static T* self(lua_State* L, int idx)
	{
		return (T*)luaL_checkudata(L, idx, T::meta_name);
	}
	static int gc(lua_State* L)
	{
		T* me = self(L, 1);
		me->~T();
		return 0;
	}
	int __index(lua_State* L)
	{
		if (lua_type(L, 2) == LUA_TSTRING) {
			const char* k = lua_tostring(L, 2);
			const luaL_Reg* func = T::instance_funcs;
			while (func->name) {
				if (!strcmp(func->name, k)) {
					lua_pushcfunction(L, func->func);
					return 1;
				}
				func++;
			}
		}
		return getItem(L);
	}
	int __newindex(lua_State* L)
	{
		return setItem(L);
	}
	virtual int getItem(lua_State* L) { return 0; };
	virtual int setItem(lua_State* L) { return 0; };
};

struct CBuffer : public LuaObject<CBuffer> {
	char* buffer;
	size_t size;
	bool owner;
	lua_State* L;
	int dbg_ref_count;
	CBuffer(lua_State* _L)
		: L(_L)
	{
		buffer = 0;
		size = 0;
		owner = false;
		dbg_ref_count = 0;
	}
	virtual ~CBuffer()
	{
		if (owner)
			free(buffer);
		// printf("CBuffer::refcount=%d\n",dbg_ref_count);
	}
	void release()
	{
		if (owner) {
			free(buffer);
		}
		buffer = 0;
		size = 0;
		owner = false;
	}
	int alloc(lua_State* L)
	{
		release();
		size = lua_tointeger(L, 2);
		buffer = (char*)malloc(size);
		if (!buffer)
			luaL_error(L, "alloc buffer failed");
		owner = true;
		return 0;
	}
	int bind(lua_State* L)
	{
		CBuffer* other = self(L, 2);
		release();
		buffer = other->buffer;
		size = other->size;
		owner = false;
		return 0;
	}
	int load(lua_State* L)
	{
		release();
		const char* p = luaL_checklstring(L, 2, &size);
		if (!size)
			luaL_error(L, "Invalid buffer in CBuffer::load().");
		buffer = (char*)malloc(size);
		if (!buffer) {
			luaL_error(L, "alloc buffer failed.");
		}
		memcpy(buffer, p, size);
		owner = true;
		return 0;
	}
	int toString(lua_State* L)
	{
		lua_pushlstring(L, buffer, size);
		return 1;
	}
	int __len(lua_State* L)
	{
		lua_pushinteger(L, size);
		return 1;
	}
	operator char*() const
	{
		return buffer;
	}
	static constexpr const char* meta_name = "gamex::CBuffer";
	static constexpr const char* class_name = "CBuffer";
	static const luaL_Reg instance_funcs[];
	static const luaL_Reg meta_funcs[];
	static const luaL_Reg static_funcs[];
};

const luaL_Reg CBuffer::instance_funcs[] = {
	INST_FUNC(CBuffer, bind),
	INST_FUNC(CBuffer, load),
	INST_FUNC(CBuffer, alloc),
	INST_FUNC(CBuffer, toString),
	{ 0, 0 }
};
const luaL_Reg CBuffer::meta_funcs[] = {
	INST_FUNC(CBuffer, __index),
	INST_FUNC(CBuffer, __newindex),
	INST_FUNC(CBuffer, __len),
	{ 0, 0 }
};
const luaL_Reg CBuffer::static_funcs[] = {
	STATIC_FUNC(CBuffer, create),
	{ 0, 0 }
};

struct CAccessor : public LuaObject<CAccessor> {
	rtk::LuaTableRef ref_type;
	int ref_buffer;
	CBuffer* buffer_obj;
	int buffer_offset;
	lua_State* L;
	int arr_len;
	size_t type_size;
	std::string type_name;
	std::string char_encoding;
	using CObjectCache = std::map<int, rtk::LuaTableRef>;
	CObjectCache cached_fields;

	CAccessor(lua_State* _L);
	virtual ~CAccessor();
	int __len(lua_State* L);
	int toString(lua_State* L)
	{
		lua_pushlstring(L, getBuffer(), type_size);
		return 1;
	}
	virtual int getItem(lua_State* L) override;
	virtual int setItem(lua_State* L) override;
	int toTable(lua_State* L);
	int bind(lua_State* L);
	int setEncoding(lua_State* L);
	int getEncoding(lua_State* L);
	int arrayGetItem(lua_State* L);
	int structGetItem(lua_State* L);
	int getPropertySheet(lua_State* L);
	int copy(lua_State* L);
	int memset(lua_State* L);
	int size(lua_State* L);
	int modifyItem(int offset, size_t size, const rtk::LuaVariant& typeDesc);
	int pushValue(int offet, size_t size, const rtk::LuaVariant& typeDesc);
	char* getBuffer() const;
	int doBind(int stk, CBuffer* nbuffer, size_t offset);
	void doUnbind();
	static const char* meta_name;
	static const char* class_name;
	static const luaL_Reg instance_funcs[];
	static const luaL_Reg meta_funcs[];
	static const luaL_Reg static_funcs[];
};

const char* CAccessor::meta_name = "gamex::CAccessor";
const char* CAccessor::class_name = "CAccessor";
// 实现元表
const luaL_Reg CAccessor::meta_funcs[] = {
	INST_FUNC(CAccessor, __index),
	INST_FUNC(CAccessor, __newindex),
	{ "__gc", CAccessor::gc },
	INST_FUNC(CAccessor, __len),
	{ nullptr, nullptr }
};
const luaL_Reg CAccessor::instance_funcs[] = {
	INST_FUNC(CAccessor, bind),
	INST_FUNC(CAccessor, toTable),
	INST_FUNC(CAccessor, getPropertySheet),
	INST_FUNC(CAccessor, toString),
	INST_FUNC(CAccessor, getEncoding),
	INST_FUNC(CAccessor, setEncoding),
	INST_FUNC(CAccessor, copy),
	INST_FUNC(CAccessor, memset),
	INST_FUNC(CAccessor, size),
	{ nullptr, nullptr }
};
const luaL_Reg CAccessor::static_funcs[] = {
	STATIC_FUNC(CAccessor, create),
	{ nullptr, nullptr }
};

CAccessor::CAccessor(lua_State* _L)
	: ref_type(_L, 1)
	, L(_L)
{
	buffer_obj = 0;
	ref_buffer = LUA_REFNIL;
	buffer_offset = 0;
	if (lua_type(L, 1) != LUA_TTABLE) {
		luaL_error(L, "Type is not a table %s", __FUNCTION__);
	}
	rtk::LuaVariant val = ref_type.getField("type");
	const std::string stp = val.toString();
	if (stp == "Array") {
		arr_len = (int)ref_type.getField("size");
		if (arr_len <= 0)
			luaL_error(L, "%s:%d, invalid object size %d.", __FILE__, __LINE__, arr_len);
	} else {
		arr_len = 0;
	}
	type_size = (int)ref_type.getField("FSize");
	type_name = ref_type.getField("name").toString();
	// Note : new object itself is not top of stack, so ...
	if (lua_gettop(L) > 2) {
		CBuffer* buf = CBuffer::self(L, 2);
		doBind(2, buf, (size_t)lua_tointeger(L, 3));
	}
}

CAccessor::~CAccessor()
{
	doUnbind();
}

int CAccessor::__len(lua_State* L)
{
	lua_pushinteger(L, arr_len);
	return 1;
}

int CAccessor::setEncoding(lua_State* L)
{
	const char* p = lua_tostring(L, 2);
	if (!p)
		p = "UTF-8";
	char_encoding = p;
	return 0;
}

int CAccessor::getEncoding(lua_State* L)
{
	lua_pushlstring(L, char_encoding.data(), char_encoding.size());
	return 1;
}

char* CAccessor::getBuffer() const
{
	return buffer_obj->buffer + buffer_offset;
}

int CAccessor::pushValue(int offset, size_t size, const rtk::LuaVariant& typeDesc)
{
	char* buffer = getBuffer();
	std::string stp;
	const rtk::LuaVariant bt = typeDesc["base_type"];
	if (bt.isNull())
		stp = typeDesc["type"].toString();
	else
		stp = bt.toString();
	if (stp == "Struct" || stp == "Array") {
		// 复用之前访问过的对象
		CObjectCache::iterator it = cached_fields.find(offset);
		if (it == cached_fields.end()) {
			lua_pushcfunction(L, create);
			typeDesc.toScript(L);
			lua_getref(L, ref_buffer);
			lua_pushinteger(L, buffer_offset + offset);
			lua_call(L, 3, 1);
			cached_fields.insert({ offset, rtk::LuaTableRef(L, -1) });
		} else {
			// utils_debug("Using cached obj at 0x%x of type %s %s\n", offset, stp.c_str(), typeDesc["name"].toString().c_str());
			it->second.getref();
		}
		return 1;
	}
	if (stp == "int64") {
		int64_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "uint64") {
		uint64_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "int32") {
		int32_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "uint32") {
		uint32_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "int16") {
		int16_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "uint16") {
		uint16_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "int8") {
		int8_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
	} else if (stp == "uint8") {
		uint8_t val;
		memcpy(&val, buffer + offset, sizeof(val));
		lua_pushinteger(L, val);
		return 1;
	} else if (stp == "char") {
		lua_pushlstring(L, buffer + offset, 1);
	} else if (stp == "hex") {
		std::string val;
		int i;
		for (i = 0; i < size; i++) {
			uint8_t t = buffer[offset + i];
			char cc[3];
			sprintf(cc, "%02x", t);
			val += cc;
		}
		lua_pushlstring(L, val.c_str(), val.size());
	} else if (stp == "string") {
		if (char_encoding.size() && char_encoding != "UTF-8") {
			std::string is(buffer + offset, size);
			is = rtk_iconv(char_encoding.c_str(),
				"UTF-8", is.c_str(), is.size());
			lua_pushlstring(L, is.data(), is.size());
		} else {
			lua_pushlstring(L, buffer + offset, size);
		}
	} else if (stp == "bool") {
		uint8_t val = buffer[offset];
		lua_pushboolean(L, val ? 1 : 0);
	} else {
		luaL_error(L, "%s:%d, invalid type spec.", __FILE__, __LINE__);
	}
	return 1;
}

int CAccessor::arrayGetItem(lua_State* L)
{
	int key = lua_tointeger(L, 2);
	if (key <= 0 || key > arr_len)
		luaL_error(L, "%s:%d subscript out of range.", __FILE__, __LINE__);
	rtk::LuaVariant ele_type = ref_type.getField("ele_type");
	int elementSize = (int)ref_type.getField("FElementSize");
	int offset = (key - 1) * elementSize;
	if ((size_t)offset >= type_size)
		luaL_error(L, "%s:%d subscript out of range.", __FILE__, __LINE__);
	pushValue(offset, elementSize, ele_type);
	return 1;
}

int CAccessor::structGetItem(lua_State* L)
{
	const char* key = lua_tostring(L, 2);
	if (!key)
		return 0;
	rtk::LuaVariant typeDesc;
	typeDesc = ref_type.getField("items")[key];
	if (!typeDesc.isTable()) {
		luaL_error(L, "invalid item spec of '%s.%s'.", type_name.c_str(), key);
	}
	int offset = (int)typeDesc["FOffset"];
	int size = (int)typeDesc["FSize"];
	pushValue(offset, size, typeDesc["type"]);
	return 1;
}

int CAccessor::getItem(lua_State* L)
{
	const char* k = lua_tostring(L, 2);
	if (k) {
		if (!strcmp(k, "FDef")) {
			// get type definition
			ref_type.getref();
			return 1;
		} else if (!strcmp(k, "buffer")) {
			if (buffer_obj) {
				lua_getref(L, ref_buffer);
				return 1;
			} else {
				return 0;
			}
		}
	}
	if (!buffer_obj) {
		luaL_error(L, "%s : object not bound.", __FUNCTION__);
	}
	if (arr_len > 0) {
		return arrayGetItem(L);
	} else {
		return structGetItem(L);
	}
}

int CAccessor::modifyItem(int offset, size_t size, const rtk::LuaVariant& typeDesc)
{
	char* buffer = getBuffer();
	int vtype = lua_type(L, 3);
	std::string stp;
	const rtk::LuaVariant vtp = typeDesc["base_type"];
	if (vtp.isNull())
		stp = typeDesc["type"].toString();
	else
		stp = vtp.toString();
	if (stp == "Struct" || stp == "Array") {
		// 只有buffer和userdata可赋值给结构体字段
		if (vtype == LUA_TSTRING) {
			if (lua_len(L, 3) != size) {
				luaL_error(L, "%s:%d, size mismatch in setItem.", __FILE__, __LINE__);
			}
			memcpy(buffer + offset, lua_tostring(L, 3), size);
			return 0;
		} else if (vtype == LUA_TUSERDATA) {
			CAccessor* oval = (CAccessor*)luaL_checkudata(L, 3, meta_name);
			if (oval->type_size != size) {
				luaL_error(L, "%s:%d, size mismatch in setItem.", __FILE__, __LINE__);
			}
			memcpy(buffer + offset, oval->getBuffer(), size);
			return 0;
		}
		luaL_error(L, "%s:%d, Invalid value", __FILE__, __LINE__);
	}
	// basic types
	if (stp == "int64") {
		int64_t val = lua_tointeger(L, 3);
		memcpy(buffer + offset, &val, sizeof(val));
	} else if (stp == "uint64") {
		uint64_t val = lua_tointeger(L, 3);
		memcpy(buffer + offset, &val, sizeof(val));
	} else if (stp == "int32") {
		int32_t val = lua_tointeger(L, 3);
		memcpy(buffer + offset, &val, sizeof(val));
	} else if (stp == "uint32") {
		uint32_t val = lua_tointeger(L, 3);
		memcpy(buffer + offset, &val, sizeof(val));
	} else if (stp == "int16") {
		int16_t val = lua_tointeger(L, 3);
		memcpy(buffer + offset, &val, sizeof(val));
	} else if (stp == "uint16") {
		uint16_t val = lua_tointeger(L, 3);
		memcpy(buffer + offset, &val, sizeof(val));
	} else if (stp == "int8") {
		buffer[offset] = (int8_t)lua_tointeger(L, 3);
	} else if (stp == "uint8") {
		buffer[offset] = (uint8_t)lua_tointeger(L, 3);
	} else if (stp == "char") {
		if (vtype == LUA_TSTRING) {
			const char* vstr = lua_tostring(L, 3);
			buffer[offset] = *vstr;
		} else if (vtype == LUA_TNUMBER) {
			buffer[offset] = lua_tointeger(L, 3);
		} else {
			// incompatible type
			luaL_error(L, "%s:%d incompatible type for char.", __FILE__, __LINE__);
		}
	} else if (stp == "hex") {
		int i;
		if (vtype != LUA_TSTRING || lua_len(L, 3) != size * 2)
			luaL_error(L, "%s:%d incompatible type for hex.", __FILE__, __LINE__);
		const char* src = lua_tostring(L, 3);
		char* tgt = buffer + offset;
		for (i = 0; i < size; i += 2, tgt++, src += 2) {
			uint8_t a = src[i];
			uint8_t b = src[i + 1];
			if (!isxdigit(a) || !isxdigit(b))
				luaL_error(L, "%s:%d incompatible type for hex.", __FILE__, __LINE__);
			auto fromHex = [](uint8_t val) -> uint8_t {
				if (val >= '0' && val <= '9')
					return val - '0';
				if (val >= 'a' && val <= 'f')
					return val - 'a' + 0x10;
				if (val >= 'F' && val <= 'F')
					return val - 'A' + 0x10;
				assert(0);
				return 0;
			};
			*tgt = (fromHex(a) << 4) | fromHex(b);
		}
	} else if (stp == "string") {
		if (vtype != LUA_TSTRING) {
			// incompatible type
			luaL_error(L, "%s:%d incompatible type for string.", __FILE__, __LINE__);
		}
		if (char_encoding.size() && char_encoding != "UTF-8") {
			size_t len;
			const char* p = luaL_checklstring(L, 3, &len);
			std::string is = rtk_iconv("UTF-8", char_encoding.c_str(),
				p, len);
			if (is.size() > size)
				luaL_error(L, "%s:%d string too long.", __FILE__, __LINE__);
			::memset(buffer + offset, 0, size);
			memcpy(buffer + offset, is.data(), is.size());
		} else {
			size_t len = lua_len(L, 3);
			if (len > size)
				luaL_error(L, "%s:%d string too long.", __FILE__, __LINE__);
			::memset(buffer + offset, 0, size);
			memcpy(buffer + offset, luaL_checkstring(L, 3), len);
		}
	} else if (stp == "bool") {
		buffer[offset] = lua_toboolean(L, 3) ? 1 : 0;
	} else {
		luaL_error(L, "%s:%d, invalid type spec.", __FILE__, __LINE__);
	}
	return 0;
}

int CAccessor::setItem(lua_State* L)
{
	if (!buffer_obj) {
		luaL_error(L, "%s : object not bound.", __FUNCTION__);
	}
	rtk::LuaVariant typeDesc;
	if (arr_len > 0) {
		int key = lua_tointeger(L, 2);
		if (key <= 0 || key > arr_len)
			luaL_error(L, "%s:%d subscript out of range.", __FILE__, __LINE__);
		rtk::LuaVariant ele_type = ref_type.getField("ele_type");
		int elementSize = (int)ref_type.getField("FElementSize");
		int offset = (key - 1) * elementSize;
		if ((size_t)offset >= type_size)
			luaL_error(L, "%s:%d subscript out of range.", __FILE__, __LINE__);
		modifyItem(offset, elementSize, ele_type);
	} else {
		const char* key = lua_tostring(L, 2);
		if (!key)
			return 0;
		typeDesc = ref_type.getField("items")[key];
		int offset = (int)typeDesc["FOffset"];
		int size = (int)typeDesc["FSize"];
		modifyItem(offset, size, typeDesc["type"]);
	}
	return 0;
}

int CAccessor::toTable(lua_State* L)
{
	luartk_require_with_ret(L, RTK_MODULE_NAME);
	lua_getfield(L, -1, "__obj_to_table");
	lua_remove(L, -2);
	lua_pushvalue(L, 1);
	lua_call(L, 1, 1);
	return 1;
}

// obj:getPropertySheet(Builder)
int CAccessor::getPropertySheet(lua_State* L)
{
	luartk_require_with_ret(L, RTK_MODULE_NAME);
	lua_getfield(L, -1, "__obj_to_prop_sheet");
	lua_remove(L, -2);
	lua_insert(L, 1);
	lua_call(L, lua_gettop(L) - 1, 1);
	return 1;
}

int CAccessor::copy(lua_State* L)
{
	if (!buffer_obj) {
		luaL_error(L, "%s : object not bound.", __FUNCTION__);
	}
	int tp = lua_type(L, 2);
	const char* src = nullptr;
	if (tp == LUA_TSTRING) {
		size_t got_size;
		src = lua_tolstring(L, 2, &got_size);
		if (got_size != type_size) {
			luaL_error(L, "CAccessor::copy: size mismatch, expected %d, got %d.",
				type_size, got_size);
		}
	} else if (tp == LUA_TUSERDATA) {
		lua_getmetatable(L, 2);
		luaL_getmetatable(L, CAccessor::meta_name);
		if (lua_rawequal(L, -1, -2)) {
			lua_pop(L, 2);
			CAccessor* other = CAccessor::self(L, 2);
			if (other->type_size != type_size) {
				luaL_error(L, "CAccessor::copy: size mismatch, expected %d, got %d.",
					type_size, other->type_size);
			}
			src = other->buffer_obj->buffer + other->buffer_offset;
		} else {
			lua_pop(L, 2);
		}
	}
	if (!src) {
		luaL_error(L, "CAccessor::copy() : invalid type.");
		return 0;
	}
	::memcpy(buffer_obj->buffer + buffer_offset, src, type_size);
	return 0;
}

int CAccessor::memset(lua_State* L)
{
	if (!buffer_obj) {
		luaL_error(L, "%s : object not bound.", __FUNCTION__);
	}
	int tp = lua_type(L, 2);
	char ch;
	if (tp == LUA_TNUMBER)
		ch = lua_tointeger(L, 2);
	else if (tp == LUA_TSTRING) {
		if (lua_len(L, 2) != 1)
			luaL_error(L, "memset(ch):ch should be int or char");
		const char* s = lua_tostring(L, 2);
		ch = *s;
	}
	::memset(buffer_obj->buffer + buffer_offset, ch, type_size);
	return 0;
}

int CAccessor::size(lua_State* L)
{
	lua_pushinteger(L, type_size);
	return 1;
}

/*
	obj:bind(buffer,offset)
*/
int CAccessor::doBind(int stkIdx, CBuffer* nbuffer, size_t offset)
{
	if (offset + type_size > nbuffer->size) {
		luaL_error(L, "buffer not big enough(offset=%d,size=%d,buffersize=%d).", offset, type_size, nbuffer->size);
	}
	doUnbind();
	buffer_obj = nbuffer;
	buffer_offset = offset;
	lua_pushvalue(L, stkIdx);
	ref_buffer = luaL_ref(L, LUA_REGISTRYINDEX);
	buffer_obj->dbg_ref_count++;
	return 0;
}

int CAccessor::bind(lua_State* L)
{
	CBuffer* nbuffer = CBuffer::self(L, 2);
	size_t offset = lua_tointeger(L, 3);
	doBind(2, nbuffer, offset);
	return 0;
}

void CAccessor::doUnbind()
{
	if (buffer_obj) {
		assert(ref_buffer != LUA_REFNIL);
		lua_unref(L, ref_buffer);
		buffer_obj->dbg_ref_count--;
		buffer_obj = 0;
		ref_buffer = LUA_REFNIL;
	}
}

static int luaopen_CAccessor(lua_State* L)
{
	CAccessor::Register(L);
	CBuffer::Register(L);
	return 0;
}

GAMEX_ADD_INIT_FUNC(luaopen_CAccessor)
